package com.zhxt.common.filter;

import com.zhxt.common.enums.HttpMethod;
import com.zhxt.common.utils.StringUtils;
import org.apache.commons.lang3.StringEscapeUtils;

import java.io.*;
import java.util.*;
import javax.servlet.*;
import javax.servlet.annotation.WebFilter;
import javax.servlet.http.*;


@WebFilter(urlPatterns = "/*")
public class WebXssFilter implements Filter {
    /**
     * 排除链接
     */
    public List<String> excludes = new ArrayList<>();

    @Override
    public void init(FilterConfig filterConfig) throws ServletException {
        // 初始化过滤器
        String tempExcludes = filterConfig.getInitParameter("excludes");
        if (StringUtils.isNotEmpty(tempExcludes)) {
            String[] url = tempExcludes.split(",");
            for (int i = 0; url != null && i < url.length; i++) {
                excludes.add(url[i]);
            }
        }
    }

    @Override
    public void doFilter(ServletRequest request, ServletResponse response, FilterChain chain)
            throws IOException, ServletException {
        // 创建自定义的请求包装器，对请求参数进行XSS过滤
        HttpServletRequest wrappedRequest = new CustomHttpServletRequest((HttpServletRequest) request);
        // 包装响应对象，对响应内容进行XSS过滤
        HttpServletResponse wrappedResponse = new CustomHttpServletResponse((HttpServletResponse) response);
        if (handleExcludeURL(wrappedRequest, wrappedResponse)) {
            chain.doFilter(wrappedRequest, wrappedResponse);
            return;
        }
        XssHttpServletRequestWrapper xssRequest = new XssHttpServletRequestWrapper((HttpServletRequest) request);

        // 将包装后的请求和响应对象传递给下一个过滤器或目标资源
        chain.doFilter(xssRequest, wrappedResponse);
    }
    private boolean handleExcludeURL(HttpServletRequest request, HttpServletResponse response)
    {
        String url = request.getServletPath();
        String method = request.getMethod();
        // GET DELETE 不过滤
        if (method == null || HttpMethod.GET.matches(method) || HttpMethod.DELETE.matches(method))
        {
            return true;
        }
        return StringUtils.matches(url, excludes);
    }
    @Override
    public void destroy() {
        // 销毁过滤器
    }

    // 自定义的HttpServletRequest包装器，用于对请求参数进行XSS过滤
    public static class CustomHttpServletRequest extends HttpServletRequestWrapper {
        private final HttpServletRequest request;

        public CustomHttpServletRequest(HttpServletRequest request) {
            super(request);
            this.request = request;
        }

        @Override
        public String getParameter(String name) {
            String value = request.getParameter(name);
            return cleanXSS(value);
        }

        @Override
        public String[] getParameterValues(String name) {
            String[] values = request.getParameterValues(name);
            if (values != null) {
                for (int i = 0; i < values.length; i++) {
                    values[i] = cleanXSS(values[i]);
                }
            }
            return values;
        }

        @Override
        public Map<String, String[]> getParameterMap() {
            Map<String, String[]> parameterMap = request.getParameterMap();
            Map<String, String[]> cleanedMap = new HashMap<>();
            for (Map.Entry<String, String[]> entry : parameterMap.entrySet()) {
                String[] values = entry.getValue();
                if (values != null) {
                    for (int i = 0; i < values.length; i++) {
                        values[i] = cleanXSS(values[i]);
                    }
                }
                cleanedMap.put(entry.getKey(), values);
            }
            return cleanedMap;
        }

        // 对输入进行XSS过滤的方法
        private String cleanXSS(String value) {
            if (value == null) {
                return null;
            }
            // 使用Apache Commons Text的StringEscapeUtils进行HTML编码，移除潜在的XSS攻击代码
            return StringEscapeUtils.escapeHtml4(value);
        }
    }

    // 自定义的HttpServletResponse包装器，用于对响应内容进行XSS过滤
    public static class CustomHttpServletResponse extends HttpServletResponseWrapper {
        private final HttpServletResponse response;
        private PrintWriter writer;

        public CustomHttpServletResponse(HttpServletResponse response) {
            super(response);
            this.response = response;
        }

        @Override
        public PrintWriter getWriter() throws IOException {
            if (writer == null) {
                writer = new PrintWriter(new OutputStreamWriter(new XSSFilterOutputStream(response.getOutputStream()), "UTF-8"));
            }
            return writer;
        }
    }

    // 自定义的OutputStream，用于对响应内容进行XSS过滤
    public static class XSSFilterOutputStream extends OutputStream {
        private final OutputStream outputStream;
        private final ByteArrayOutputStream buffer = new ByteArrayOutputStream();

        public XSSFilterOutputStream(OutputStream outputStream) {
            this.outputStream = outputStream;
        }

        @Override
        public void write(int b) throws IOException {
            buffer.write(b);
        }

        @Override
        public void flush() throws IOException {
            // 对缓冲区中的内容进行XSS过滤
            String content = buffer.toString("UTF-8");
            String filteredContent = cleanXSS(content);
            outputStream.write(filteredContent.getBytes("UTF-8"));
            outputStream.flush();
            buffer.reset();
        }

        @Override
        public void close() throws IOException {
            flush();
            outputStream.close();
        }

        // 对响应内容进行XSS过滤的方法
        private String cleanXSS(String content) {
            if (content == null) {
                return null;
            }
            // 使用Apache Commons Text的StringEscapeUtils进行HTML编码，移除潜在的XSS攻击代码
            return StringEscapeUtils.escapeHtml4(content);
        }
    }
}